{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%reload_ext autoreload\n",
    "%autoreload 2\n",
    "\n",
    "import sys\n",
    "sys.path.append(os.path.abspath('../'))\n",
    "from common import *\n",
    "import configs\n",
    "\n",
    "fig_name = 'normalization'\n",
    "fig_dir = join(FIGURE_DIR, fig_name)\n",
    "used_configs = [configs.Warp(), configs.WarpNotNormalized(), configs.FiniteDifferences()]\n",
    "names = ['Reference (FD)', 'No normalization', 'With normalization']\n",
    "\n",
    "sdf_resolutions = [64, 128, 256]\n",
    "y_offset = -0.2\n",
    "fontsize = 12\n",
    "base_size = 3\n",
    "n_rows = 3\n",
    "n_cols = 3\n",
    "total_width = base_size * 3\n",
    "fig = plt.figure(1, figsize=(COLUMN_WIDTH, 0.78 * COLUMN_WIDTH), constrained_layout=False)\n",
    "gs = fig.add_gridspec(n_rows, n_cols, wspace=0.025, hspace=0.025)\n",
    "r = 80\n",
    "for row, sdf_res in enumerate(sdf_resolutions):\n",
    "    for col, (config, name) in enumerate(zip(used_configs, names)):\n",
    "        ax = fig.add_subplot(gs[row, col])\n",
    "        if col == 1 and (row == 0 or row == 1):\n",
    "            ax.arrow(0.4, 0.07, 0.02, 0.09, head_width=0.02, length_includes_head=True, facecolor='k')\n",
    "        if col == 1 and row == 1:\n",
    "            ax.arrow(0.82, 0.685, -0.08, -0.02, head_width=0.02, length_includes_head=True, facecolor='k')\n",
    "        if col == 1 and row == 0:\n",
    "            ax.arrow(0.75, 0.7, -0.05, -0.07, head_width=0.02, length_includes_head=True, facecolor='k')\n",
    "            ax.arrow(0.99, 0.45, -0.1, -0.02, head_width=0.02, length_includes_head=True, facecolor='k')\n",
    "        if col == 0:\n",
    "            ax.set_ylabel(f'${sdf_res}^3$', labelpad=5, fontsize=fontsize)\n",
    "        if col == 0 and row == 1:\n",
    "            ax.text(-0.2, 0.4, 'Grid resolution (whole shape)', rotation=90, va='center',\n",
    "                    ha='center', fontsize=fontsize)\n",
    "        img = read_img(join(fig_dir, f'{config.name}_{sdf_res}_grad.exr'), tonemap=False)\n",
    "        img = img[160:-100, 40:-80, :]\n",
    "        aspect = img.shape[0] / img.shape[1]\n",
    "        ax.imshow(np.mean(img, axis=-1), cmap='coolwarm_r', vmin=-r, vmax=r,\n",
    "                  interpolation='none', extent=[0, 1, 0, aspect])\n",
    "        disable_ticks(ax)\n",
    "        if row == n_rows - 1:\n",
    "            txt = ax.set_title(name, fontsize=fontsize, y=y_offset)\n",
    "\n",
    "plt.margins(0, 0)\n",
    "# save_fig(fig_name)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3.10.4 ('mi')",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.4"
  },
  "orig_nbformat": 4,
  "vscode": {
   "interpreter": {
    "hash": "c8cc192903caa17681aa39d71202092ac11a526d37a1c4ad2948f13605924304"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
